package com.dimples.dd.gateway.filter.defend;

import cn.hutool.core.codec.Base64Encoder;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSONObject;
import com.dimples.dd.common.constant.DDConstant;
import com.dimples.dd.common.result.CommonResult;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.util.AntPathMatcher;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.List;

import static com.dimples.dd.common.result.ResultCode.URL_FORBID_ERROR;

/**
 * 微服务防护
 * <p>
 * 1、拦截所有网关的请求，在请求头中加入网关密钥
 * <p>
 * 2、在微服务模块验证请求头中是否带了网关密钥，如果没有，不予访问
 *
 * @author zhongyj <1126834403@qq.com><br/>
 * @date 2024/6/7
 */
@Slf4j
@Component
public class DDGatewayRequestFilter implements GlobalFilter {

    @Value("${dd.gateway.forbid-request-uri:}")
    private String forbidRequestUri;

    private final AntPathMatcher pathMatcher = new AntPathMatcher();

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        ServerHttpResponse response = exchange.getResponse();

        // 禁止客户端的访问资源逻辑
        Mono<Void> checkForbidUriResult = checkForbidUri(request, response);
        if (checkForbidUriResult != null) {
            return checkForbidUriResult;
        }

        byte[] token = Base64Encoder.encode(DDConstant.GATEWAY_TOKEN_VALUE.getBytes()).getBytes();
        String[] headerValues = {new String(token)};
        ServerHttpRequest build = exchange.getRequest().mutate().header(DDConstant.GATEWAY_TOKEN_HEADER, headerValues).build();
        ServerWebExchange newExchange = exchange.mutate().request(build).build();

        return chain.filter(newExchange);
    }

    private Mono<Void> checkForbidUri(ServerHttpRequest request, ServerHttpResponse response) {
        String uri = request.getPath().toString();
        boolean shouldForward = true;
        List<String> forbidRequestUris = StrUtil.split(forbidRequestUri, ",");
        if (CollUtil.isNotEmpty(forbidRequestUris)) {
            for (String u : forbidRequestUris) {
                if (pathMatcher.match(u, uri)) {
                    shouldForward = false;
                }
            }
        }
        if (!shouldForward) {
            CommonResult<String> ddResponse = CommonResult.error(URL_FORBID_ERROR);
            return makeResponse(response, ddResponse);
        }
        return null;
    }

    private Mono<Void> makeResponse(ServerHttpResponse response, CommonResult<String> ddResponse) {
        response.setStatusCode(HttpStatus.FORBIDDEN);
        response.getHeaders().add(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE);
        DataBuffer dataBuffer = response.bufferFactory().wrap(JSONObject.toJSONString(ddResponse).getBytes());
        return response.writeWith(Mono.just(dataBuffer));
    }
}
